import numpy as np
import matplotlib.pyplot as plt


class ServoMotor:
    def __init__(self, k_m, J, B, tau_ext, time_constant=0.1):
        """
        k_m: 力矩常数 (N·m/A)
        J: 电机转动惯量 (kg·m²)
        B: 电机粘性阻尼系数 (N·m·s)
        tau_ext: 外部扰动力矩 (N·m)
        time_constant: 电机的一阶延迟时间常数 (s)
        """
        self.k_m = k_m  # 力矩常数
        self.J = J  # 转动惯量
        self.B = B  # 粘性阻尼
        self.tau_ext = tau_ext  # 外部扰动力矩
        self.time_constant = time_constant  # 一阶延迟时间常数

        self.theta = 0  # 电机角度 (rad)
        self.omega = 0  # 电机角速度 (rad/s)
        self.alpha = 0  # 电机角加速度 (rad/s²)
        self.I = 0  # 电机输入电流 (A)

    def update(self, T_motor, dt):
        """
        根据电流更新电机状态
        I: 电机输入电流 (A)
        dt: 时间步长 (s)
        """
        # 计算电机的力矩
        # T_motor = self.k_m * I  # 力矩与电流的关系

        # 电机动力学方程：T_motor = J * alpha + B * omega + T_ext
        # 角加速度：alpha = (T_motor - B * omega - T_ext) / J
        self.alpha = (T_motor - self.B * self.omega - self.tau_ext) / self.J

        # 更新角速度和角度（考虑一阶延迟环节）
        self.omega += self.alpha * dt
        self.theta += self.omega * dt

        # 添加一阶惯性环节（电机的动态响应）
        self.omega = self.omega * np.exp(-dt / self.time_constant)

        return self.theta, self.omega, self.alpha


class Joint:
    def __init__(self, motor, gear_ratio, K_joint):
        """
        motor: 对应的电机对象
        gear_ratio: 电机与关节的减速比
        K_joint: 关节的刚度常数
        """
        self.motor = motor  # 关节对应的电机
        self.gear_ratio = gear_ratio  # 减速比
        self.K_joint = K_joint  # 关节刚度常数
        self.joint_angle = 0  # 关节角度 (rad)
        self.joint_omega = 0  # 关节角速度 (rad/s)
        self.joint_alpha = 0  # 关节角加速度 (rad/s²)

    def set_joint_angle(self, target_angle, dt):
        """
        设置目标关节角度，通过目标角度计算期望力矩，并更新电机状态
        target_angle: 目标关节角度 (rad)
        dt: 时间步长 (s)
        """
        # 计算目标力矩 (T_joint) = K_joint * target_angle
        target_torque = self.K_joint * target_angle

        # 通过减速比转换力矩
        motor_torque = target_torque

        # 计算电机输入电流 (I) = T_motor / k_m
        I = motor_torque / self.motor.k_m

        # 更新电机状态
        motor_angle, motor_omega, motor_alpha = self.motor.update(I, dt)

        # 通过减速比转换电机角度和关节角度
        self.joint_angle = motor_angle / self.gear_ratio
        self.joint_omega = motor_omega / self.gear_ratio
        self.joint_alpha = motor_alpha / self.gear_ratio

        return self.joint_angle, self.joint_omega, self.joint_alpha


# 电机参数
k_m = 0.1  # 力矩常数 (N·m/A)
J = 0.01  # 转动惯量 (kg·m²)
B = 0.001  # 粘性阻尼系数 (N·m·s)
tau_ext = 0.0  # 外部扰动力矩 (N·m)
time_constant = 0.1  # 电机的一阶延迟时间常数 (s)

# 初始化电机和关节
motor = ServoMotor(k_m, J, B, tau_ext, time_constant)
gear_ratio = 10  # 电机与关节的减速比
K_joint = 0.5  # 关节刚度常数
joint = Joint(motor, gear_ratio, K_joint)

# 控制参数
dt = 0.01  # 时间步长 (s)
num_steps = 1000  # 总的仿真步数

# 记录结果
joint_angles = []
joint_omegas = []
joint_alphas = []

for step in range(num_steps):
    target_angle = np.sin(0.01 * step)  # 假设目标关节角度为正弦波
    joint_angle, joint_omega, joint_alpha = joint.set_joint_angle(target_angle, dt)

    joint_angles.append(joint_angle)
    joint_omegas.append(joint_omega)
    joint_alphas.append(joint_alpha)

# 绘制结果
time = np.arange(0, num_steps * dt, dt)
plt.figure(figsize=(12, 6))

plt.subplot(3, 1, 1)
plt.plot(time, joint_angles)
plt.title("Joint Angle vs Time")
plt.xlabel("Time (s)")
plt.ylabel("Joint Angle (rad)")

plt.subplot(3, 1, 2)
plt.plot(time, joint_omegas)
plt.title("Joint Angular Velocity vs Time")
plt.xlabel("Time (s)")
plt.ylabel("Joint Angular Velocity (rad/s)")

plt.subplot(3, 1, 3)
plt.plot(time, joint_alphas)
plt.title("Joint Angular Acceleration vs Time")
plt.xlabel("Time (s)")
plt.ylabel("Joint Angular Acceleration (rad/s²)")

plt.tight_layout()
print(joint_angles)
plt.show()

